import os
from skimage import io, img_as_float32
from skimage.color import gray2rgb
from sklearn.model_selection import train_test_split
from imageio import mimread

import numpy as np
from torch.utils.data import Dataset
import pandas as pd
from AI.augmentation import AllAugmentationTransform
import glob


def read_video(name, frame_shape):
    """
    Read video which can be:
      - an image of concatenated frames
        帧连接的图像（mp4 gif）
      - '.mp4' and'.gif'
      - folder with videos
    """
    # frame_shape输入实例 (256, 256, 3)
    # 如果name是文件夹路径（这个文件夹下必须是mp4 gif），则 io.imread 每一个mp4 或者 gif 文件
    # 数据的结构 np.array[ [io.imread("")], [io.imread("")] ]
    # 数组中的每一个都是三维数组（大小和维度与原视频有关）
    if os.path.isdir(name):
        frames = sorted(os.listdir(name))
        num_frames = len(frames)
        video_array = np.array(
            [img_as_float32(io.imread(os.path.join(name, frames[idx]))) for idx in range(num_frames)])

    # name是图片文件所在路径+名字
    # 将name所有的大写字母转化位小写，并判断是否以 .png .jpg结尾
    elif name.lower().endswith('.png') or name.lower().endswith('.jpg'):
        image = io.imread(name)
        # 图片的第三维度位 1  或  2 则转化为灰度图
        if len(image.shape) == 2 or image.shape[2] == 1:
            image = gray2rgb(image)
        # 图片的第三维度为4 则只保留前三个
        if image.shape[2] == 4:
            image = image[..., :3]

        image = img_as_float32(image)
        # 将数据的第二维度和第一维度交换 比如 np.moveaxis(np.zeros((3, 4, 5)), 1, 0) 转化后为 (4, 3, 5)
        video_array = np.moveaxis(image, 1, 0)

        # 将元组(-1, ) 和元组frame_shape 合并(-1, 256, 256, 3) 既将reshape成任意行的【256， 256， 3】
        video_array = video_array.reshape((-1,) + frame_shape)

        # 将数据的第二维度和第三维度交换 比如 np.moveaxis(np.zeros((3, 4, 5)), 1, 2) 转化后为 (3, 5, 4)
        # 这里将每个数据的行和列互换
        video_array = np.moveaxis(video_array, 1, 2)

    # name是视频文件的路径+名字
    elif name.lower().endswith('.gif') or name.lower().endswith('.mp4') or name.lower().endswith('.mov'):

        # imageio.mimread: 读取gif，每一帧会存放到list的一个位置中
        # imageio.mimsave: 保存gif，输入也是一个list数组
        video = np.array(mimread(name))

        # 视频是三维的则转化为灰度图
        if len(video.shape) == 3:
            video = np.array([gray2rgb(frame) for frame in video])
        if video.shape[-1] == 4:
            video = video[..., :3]
        video_array = img_as_float32(video)
    else:
        raise Exception("Unknown file extensions  %s" % name)

    # 返回结果为[ [一个图像的np数组], [一个图像的np数组] ]
    return video_array


class FramesDataset(Dataset):
    """
    Dataset of videos, each video can be represented as:
      - an image of concatenated frames
      - '.mp4' or '.gif'
      - folder with all frames
    """

    def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_train=True,
                 random_seed=0, pairs_list=None, augmentation_params=None):
        self.root_dir = root_dir

        # 获取视频文件名列表
        self.videos = os.listdir(root_dir)

        # 数据帧的大小
        self.frame_shape = tuple(frame_shape)

        self.pairs_list = pairs_list
        self.id_sampling = id_sampling

        # 如果root_dir下有train文件夹
        if os.path.exists(os.path.join(root_dir, 'train')):

            # 断言 当后面为False时触发异常
            assert os.path.exists(os.path.join(root_dir, 'test'))
            print("Use predefined train-test split.")

            # 精选出唯一标识符 (ID) 的一个子集，以便能够继续使用。
            # 如果是则将os.path.join(root_dir, 'train')下的所有文件名以#作为分隔点，获取分割后的第一个
            if id_sampling:
                train_videos = {os.path.basename(video).split('#')[0] for video in
                                os.listdir(os.path.join(root_dir, 'train'))}
                train_videos = list(train_videos)
            else:
                train_videos = os.listdir(os.path.join(root_dir, 'train'))
            test_videos = os.listdir(os.path.join(root_dir, 'test'))
            self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test')

        # root_dir下没有train文件夹，则将数据集进行随机分割成train， test（只是文件名字分割并没有转化数据）
        else:
            print("Use random train-test split.")
            train_videos, test_videos = train_test_split(self.videos, random_state=random_seed, test_size=0.2)

        # 当前是否为训练模式
        if is_train:
            self.videos = train_videos
        else:
            self.videos = test_videos

        self.is_train = is_train

        if self.is_train:
            # 进行图片增强
            self.transform = AllAugmentationTransform(**augmentation_params)
        else:
            self.transform = None

    def __len__(self):
        return len(self.videos)

    # 输入一系列图像的path与图像的index(组合为一张图像的详细地址)，输出字典
    def __getitem__(self, idx):
        if self.is_train and self.id_sampling:
            name = self.videos[idx]

            # 随机获取一个 root_dir 文件下的所有mp4文件中的一个
            path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4')))
        else:
            # 获取一个文件的路径
            name = self.videos[idx]
            path = os.path.join(self.root_dir, name)

        # path是文件的路径
        # 返回文件名
        video_name = os.path.basename(path)

        # 如果是训练模式并且path为目录
        if self.is_train and os.path.isdir(path):

            # 获取path目录下的所有文件名
            frames = os.listdir(path)
            num_frames = len(frames)

            # 从0---num_frames-1内取两个可以重复的数字，并排序
            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))

            # 将上一部或取得frame_idx从frames内取出两个文件路径并读取
            video_array = [img_as_float32(io.imread(os.path.join(path, frames[idx]))) for idx in frame_idx]
        else:
            # video_array--数据格式[ imread(path) ]
            video_array = read_video(path, frame_shape=self.frame_shape)

            # 经过read_video的reshape后获取有几个数据
            num_frames = len(video_array)

            # 如果是训练模式从0--num_frames-1里取2个可以重复的数字不是就取range(num_frames)
            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2)) if self.is_train else range(
                num_frames)

            # 从video_array里取两个
            video_array = video_array[frame_idx]

        if self.transform is not None:
            video_array = self.transform(video_array)

        out = {}

        # 如果是训练模式
        if self.is_train:
            # 原图片数据
            source = np.array(video_array[0], dtype='float32')
            # 驱动图片数据
            driving = np.array(video_array[1], dtype='float32')

            # 将数组的第一个轴换到第三个位置其他的亦然
            out['driving'] = driving.transpose((2, 0, 1))
            out['source'] = source.transpose((2, 0, 1))
        else:
            video = np.array(video_array, dtype='float32')
            out['video'] = video.transpose((3, 0, 1, 2))

        out['name'] = video_name

        # 训练模式返回
        # {“driving”:[np.array()], "source": np.array()}

        # 非训练模式返回
        # {"name": "文件名字"}
        return out


class DatasetRepeater(Dataset):
    """
    Pass several times over the same dataset for better i/o performance
    可以通过重复传输相同的数据集来提高 I/O 性能。但是，根据应用场景，这可能并不总是最好的方法。
    你可以考虑将数据缓存起来，以减少对多次传输的需求，并加快检索速度。您还可以尝试一次将数据预处理成较小的块，
    然后再在后续操作中使用这些子集，以尽量减少数据必须传输的次数。此外，您还可以找到将过程并行化的方法，
    因为传输单个大型数据集比同时处理多个小型数据集要花费更多时间。
    """

    def __init__(self, dataset, num_repeats=100):

        # 这个就是上面实现的类
        self.dataset = dataset
        self.num_repeats = num_repeats

    def __len__(self):
        # 返回重复读取的次数 * dataset的长度
        return self.num_repeats * self.dataset.__len__()

    def __getitem__(self, idx):
        return self.dataset[idx % self.dataset.__len__()]


class PairedDataset(Dataset):
    """
    Dataset of pairs for animation.
    用于动画的数据集对
    """

    def __init__(self, initial_dataset, number_of_pairs, seed=0):

        # 最初的数据
        self.initial_dataset = initial_dataset
        pairs_list = self.initial_dataset.pairs_list

        # 设置随机种子
        np.random.seed(seed)

        # 如果成对列表为空
        if pairs_list is None:

            # 取最小的长度为最大idx
            max_idx = min(number_of_pairs, len(initial_dataset))
            nx, ny = max_idx, max_idx

            # 生成    [[0, 0],
            #        [0, 1],
            #        [0, 2],
            #        [1, 0],
            #        [1, 1],
            #        [1, 2],
            #        [2, 0],
            #        [2, 1],
            #        [2, 2]]
            xy = np.mgrid[:nx, :ny].reshape(2, -1).T

            # 获取最小的成对数组个数
            number_of_pairs = min(xy.shape[0], number_of_pairs)

            # 从xy里取随机不重复的取number_of_pairs个 [x, y]   改过代码
            self.pairs = xy.take(np.random.choice(xy.shape[0], size=number_of_pairs, replace=False, ), axis=0)
        else:
            videos = self.initial_dataset.videos

            # index 0------len(videos), name为文件名
            name_to_index = {name: index for index, name in enumerate(videos)}

            #
            pairs = pd.read_csv(pairs_list)

            # 取掩膜 isin()方法可以同时判断数据是否与多个值相等，若与其中的某个值相等则返回True，否则则为False。
            # 返回类型是，某一列与括号给定的值，与条件比较，返回一列的True， False
            pairs = pairs[np.logical_and(pairs['source'].isin(videos), pairs['driving'].isin(videos))]

            number_of_pairs = min(pairs.shape[0], number_of_pairs)
            self.pairs = []
            self.start_frames = []
            for ind in range(number_of_pairs):
                self.pairs.append(
                    (name_to_index[pairs['driving'].iloc[ind]], name_to_index[pairs['source'].iloc[ind]]))

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        pair = self.pairs[idx]
        first = self.initial_dataset[pair[0]]
        second = self.initial_dataset[pair[1]]
        first = {'driving_' + key: value for key, value in first.items()}
        second = {'source_' + key: value for key, value in second.items()}

        return {**first, **second}
